package io.a.sql.dsl;

import com.google.common.collect.ImmutableMap;
import io.a.entity.User;
import io.a.utils.H;
import lombok.extern.slf4j.Slf4j;
import org.apache.spark.sql.*;

import java.time.LocalDateTime;
import java.util.Collections;

import static io.a.utils.Const.*;
import static org.apache.spark.sql.functions.lit;

@Slf4j
public class Dsl {

    private static final String FORMAT_JDBC = "jdbc";

    private static SparkSession session;

    private static final DataFrameReader dataFrameReader = init();

    private static DataFrameReader init() {
        session = SparkSession
                .builder()
                .master("local[*]")
                .appName(Dsl.class.getName())
                .getOrCreate();

        return session.read().format(FORMAT_JDBC)
                .option(DRIVER, H.DB.getDriver()).option(URL, H.DB.getUrl())
                .option(USER, H.DB.getUser()).option(PASSWORD, H.DB.getPassword())
                /*.option("queryTimeout", "300")*/;
    }

    private static void finish() {
        if (session != null) {
            session.stop();
        }
    }


    public static void main(String[] args) {

        final Dataset<Row> userDs = dataFrameReader.option("query", "SELECT id  AS user_id, name, age FROM user").load();
        final Dataset<Row> userRoleDs = dataFrameReader.option("query", "SELECT * FROM user_role").load();

        userDs.createOrReplaceTempView("u");
        userRoleDs.createOrReplaceTempView("ur");

        // 选择列
        Dataset<Row> name = userDs.select("name");
        name.show();

        // 创建Dataset
        User user = new User(101L, "admin", 18);
        Dataset<User> javaBeanDS = session.createDataset(
                Collections.singletonList(user),
                Encoders.bean(User.class)
        );
        javaBeanDS.show();

        Dataset<Row> inner = userDs
                .join(userRoleDs, new String[]{"user_id"}, "inner")
                .withColumns(
                        ImmutableMap.of(
                                "time",
                                lit(LocalDateTime.now())));
        inner.show();

        /* 有误
        Dataset<Row> sql = session.sql("SELECT u.* FROM u JOIN ur ON u.id=ur.user_id");
        sql.show();*/

        /*Dataset<Row> userDs2 = session.sql("SELECT * FROM u");

        userDs2.show();

        Dataset<User> userDataset = userDs2.as(Encoders.bean(User.class));

        userDataset.foreach(ud -> {
            System.out.println(ud);
        });*/

        finish();
    }
}
